# -*- coding: utf-8 -*-

import tensorflow as tf
import cv2
import numpy as np
import scipy.io
import math
import random
import os
from vgg import vgg16_network


def show_density_map(density_map, name="color_map"):
    min_value, max_value = np.min(density_map), np.max(density_map)
    img = ((density_map - min_value) / max_value) * 255
    img = img.astype(np.uint8, copy=True)
    img = cv2.resize(img, (img.shape[1]*4, img.shape[0]*4))
    color_map = cv2.applyColorMap(img, cv2.COLORMAP_JET)
    cv2.imshow(name, color_map)


"""
def generate_ground_truth(mat_path, shape, image_path=None):
    shape = [shape[0]//4, shape[1]//4]
    mat = scipy.io.loadmat(mat_path)
    locations = mat["image_info"][0][0][0][0][0]
    locations = [[location[0]/8, location[1]/8] for location in locations]
    ground_truth = np.zeros(shape=shape, dtype=np.float32)
    for location in locations:
        temp = np.zeros_like(ground_truth)
        x, y = math.floor(location[0]), math.floor(location[1])
        if x >= shape[1] or y >= shape[0]:
            continue
        temp[y, x] = 1
        temp = cv2.GaussianBlur(temp, (0, 0), 5)
        ground_truth += temp
    ground_truth *= 1000
    # print(len(locations), np.sum(ground_truth))
    if image_path:
        print(len(locations), np.sum(ground_truth))
        image = cv2.imread(image_path)
        image = cv2.resize(image, dsize=(shape[1], shape[0]))
        cv2.imshow("src", image)
        show_ground_truth(ground_truth)
        cv2.destroyAllWindows()
    return ground_truth


def read_image_and_ground_truth():
    folder = "/home/lijun/data/dataset/DataSet/ShanghaiTech/part_B_final/test_data/images"
    images_name = [name for name in os.listdir(folder) if "jpg" in name]
    for name in images_name:
        name = random.choice(images_name)
        image_path = os.path.join(folder, name)
        mat_path = image_path.replace("images/", "ground_truth/GT_").replace("jpg", "mat")
        image = cv2.imread(image_path)
        image = cv2.resize(image, dsize=(512, 384))
        cv2.imshow("image", image)
        image = image.astype(dtype=np.float32, copy=False)
        image = np.divide(image, 255)
        ground_truth = generate_ground_truth(mat_path, image.shape[0:2], image_path=None)
        image = np.expand_dims(image, axis=0)
        # yield image, ground_truth
        return image, ground_truth


# def read_image_and_ground_truth():
#     folder = "/home/lijun/Dataset/crowd_counting/part_b/train/images"
#     images_name = [name for name in os.listdir(folder) if "jpg" in name]
#     # name = random.choice(images_name)
#     name = images_name[765]
#     image_path = os.path.join(folder, name)
#     image = cv2.imread(image_path)
#     cv2.imshow("image", image)
#     image = image.astype(dtype=np.float32, copy=False)
#     image = np.divide(image, 255)
#
#     mat_path = image_path.replace("images", "ground_truth").replace("jpg", "h5")
#     ground_truth_file = h5py.File(mat_path, "r")
#     ground_truth = np.array(ground_truth_file["data"])
#     ground_truth_file.close()
#
#     image = np.expand_dims(image, axis=0)
#     return image, ground_truth


def read_image():
    folder = "/home/lijun/temp/test_2"
    images_name = [name for name in os.listdir(folder) if "jpg" in name]
    for name in images_name:
        name = images_name[5]
        print(name)
        image_path = os.path.join(folder, name)
        image = cv2.imread(image_path)
        image = cv2.resize(image, (480, 270))
        cv2.imshow("image", image)
        image = image.astype(dtype=np.float32, copy=False)
        image = np.divide(image, 255)
        image = np.expand_dims(image, axis=0)
        return image


def main():
    with tf.Graph().as_default():
        input_x = tf.placeholder(dtype=tf.float32, shape=[None, None, None, 3])
        inference_op = vgg16_network(input_x)
        saver = tf.train.Saver()
        ground_truths, differences = [], []
        with tf.Session() as session:
            saver.restore(session, "./model/back/part_b_multi_1000/vgg-21000")
            # for image, ground_truth in read_image_and_ground_truth():
            #     inference = session.run(inference_op, feed_dict={input_x: image})
            #     inference = inference[0][:, :, 0]
            #     print(np.sum(inference)/1000 - np.sum(ground_truth)/1000)
            #     ground_truths.append(float(np.sum(ground_truth) / 1000))
            #     differences.append(float(np.abs(np.sum(inference)/1000 - np.sum(ground_truth)/1000)))
            # print(sum(differences) / sum(ground_truths))
            # with open("test_result_vgg_part_b_multi_10-10000.txt", "w") as result_file:
            #     result_file.writelines(" ".join([str(e) for e in ground_truths]) + "\n")
            #     result_file.writelines(" ".join([str(e) for e in differences]) + "\n")

            image, ground_truth = read_image_and_ground_truth()
            inference = session.run(inference_op, feed_dict={input_x: image})
            inference = inference[0][:, :, 0]
            print(np.sum(inference)/1000, np.sum(ground_truth)/1000)
            show_ground_truth(inference, "inference")
            show_ground_truth(ground_truth, "ground_truth")
            cv2.waitKey()
            cv2.destroyAllWindows()

            # image = read_image()
            # inference = session.run(inference_op, feed_dict={input_x: image})
            # inference = inference[0][:, :, 0]
            # print(np.sum(inference) / 1000)
            # show_ground_truth(inference, "inference")
            # cv2.waitKey()
            # cv2.destroyAllWindows()
"""


def show_labeled_points(image_path, ground_truth_file_path):
    with open(ground_truth_file_path) as ground_truth_file:
        str_ground_truths = ground_truth_file.readlines()
    ground_truths = [[int(e) for e in line.strip().split(" ")] for line in str_ground_truths]
    image = cv2.imread(image_path)
    for ground_truth in ground_truths:
        cv2.circle(image, tuple(ground_truth), 3, (0, 0, 255), -1)
    cv2.namedWindow("show", cv2.WINDOW_NORMAL)
    cv2.imshow("show", image)
    cv2.waitKey()
    cv2.destroyWindow("show")


def read_image_and_ground_truth():
    folder = "/home/lijun/temp"
    sub_folders = [os.path.join(folder, name) for name in os.listdir(folder)
                   if os.path.isdir(os.path.join(folder, name))]
    full_images_path, full_ground_truth_files_path = [], []
    for sub_folder in sub_folders:
        images_name = [name for name in os.listdir(sub_folder) if "jpg" in name]
        ground_truths_name = [name.replace("jpg", "txt") for name in images_name]
        images_path = [os.path.join(sub_folder, name) for name in images_name]
        ground_truth_files_path = [os.path.join(sub_folder, name) for name in ground_truths_name]
        full_images_path.extend(images_path)
        full_ground_truth_files_path.extend(ground_truth_files_path)
    print(len(full_images_path))
    index = 1
    for image_path, ground_truth_file_path in zip(full_images_path, full_ground_truth_files_path):
        with open(ground_truth_file_path) as ground_truth_file:
            str_ground_truths = ground_truth_file.readlines()
        ground_truths = [[int(e) for e in line.strip().split(" ")] for line in str_ground_truths]
        # show_ground_truth(image_path, ground_truth_file_path)
        image = cv2.imread(image_path)
        image = cv2.resize(image, dsize=(image.shape[1]//4, image.shape[0]//4))
        if index == 167:
            cv2.imshow("image", image)
        image = image.astype(dtype=np.float32, copy=False)
        image = np.divide(image, 255)
        image = np.expand_dims(image, axis=0)
        index += 1
        yield image, len(ground_truths)


def main():
    read_image_and_ground_truth()
    with tf.Graph().as_default():
        input_x = tf.placeholder(dtype=tf.float32, shape=[None, None, None, 3])
        inference_op = vgg16_network(input_x)
        saver = tf.train.Saver()
        with tf.Session() as session:
            saver.restore(session, "./model/back/part_b_multi_1000/vgg-21000")
            pos_differences, neg_differences, abs_differences = [], [], []
            ground_truth_total_count = 0
            index = 1
            for image, ground_truth_count in read_image_and_ground_truth():
                inference = session.run(inference_op, feed_dict={input_x: image})
                inference = inference[0][:, :, 0]
                predict_count = int(np.sum(inference)/1000)
                difference = predict_count - ground_truth_count
                if difference > 0:
                    pos_differences.append(difference)
                else:
                    neg_differences.append(difference)
                abs_differences.append(abs(difference))
                ground_truth_total_count += ground_truth_count
                if abs(difference) >= 59:
                    print(index), print(predict_count), print(ground_truth_count)
                    show_density_map(inference)
                    cv2.waitKey()
                    cv2.destroyAllWindows()
                index += 1
            pos_differences = sorted(pos_differences)
            neg_differences = sorted(neg_differences)
            abs_differences = sorted(abs_differences)
            print("max pos difference: ", pos_differences[-1], pos_differences[0])
            print("max neg difference: ", neg_differences[0], neg_differences[-1])
            print("max abs difference: ", abs_differences[-1], abs_differences[0])
            print("ratio: ", 1 - sum(abs_differences) / ground_truth_total_count)


if __name__ == "__main__":
    main()
